import sys
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import colorsys
import random
import math

point_size = 0.5

def ncolors(num):
    rgb_colors = []
    if num < 1:
        return rgb_colors
    hls_colors = []
    i = 0
    step = 360.0 / num
    while i < 360:
        h = i
        s = 90 + random.random() * 10
        l = 50 + random.random() * 10
        _hlsc = [h / 360.0, l / 100.0, s / 100.0]
        hls_colors.append(_hlsc)
        i += step
    for hlsc in hls_colors:
        r, g, b = colorsys.hls_to_rgb(hlsc[0], hlsc[1], hlsc[2])
        rgb_colors.append((r, g, b))
    return rgb_colors

# 通过交互式命令来获取文件信息
def get_info():
    print("该程序没有严格地检查输入，请确保你的输入文件是正确")
    # 读取一整行的输入
    file_Name = input("请输入文件名")
    file_Name = file_Name.strip("\"")
    file_Name = file_Name.strip("\'")
    assert file_Name.endswith(".csv"), "文件格式不正"
    # x_min = float(input("请输入x轴显示的下限"))
    # x_max = float(input("请输入x轴显示的上限"))
    # if x_min > x_max:
    #     x_min, x_max = x_max, x_min
    # y_min = float(input("请输入y轴显示的下限"))
    # y_max = float(input("请输入y轴显示的上限"))
    # if y_min > y_max:
    #     y_min, y_max = y_max, y_min
    return file_Name

def num_to_limit(num):
    ratio = 1.0 if num > 0 else -1.0
    num = abs(num)
    while num > 10.0:
        num /= 10.0
        ratio *= 10.0
    while num < 1.0:
        num *= 10.0
        ratio /= 10.0
    num = math.ceil(num)
    return num*ratio

# 读取文件
# 1. 获取有多少组数据
# 2. 创建(x,y,group)
def read_file(file_Name: str) -> tuple:
    # 输入的文件必须满足第一行是标题，后面都是数据
    # 标题格式是Applied Voltage[V],**℃
    x = []
    y = []
    group = []
    with open(file_Name, "r") as f:
        content = f.readlines()
        assert content[0][0] == "A" and len(content) > 2, "文件格式不正"
        for line in content:
            line = line.split(",")
            if "A" in line[0]:
                for i in range(len(line)):
                    if 'A' in line[i]:
                        continue
                    group.append(line[i].strip())
                    x.append([])
                    y.append([])
            else:
                for i in range(0, len(line), 2):
                    new_x = float(line[i])
                    new_y = float(line[i+1])
                    x[int(i/2)].append(new_x)
                    y[int(i/2)].append(new_y)

    return x, y, group


# 绘制图像
def draw_plot(
    x: list, 
    y: list, 
    group: list, 
    point_size: float = 20,
    title: str = None,
    label_x: str = "Applied Voltage[V]",
    label_y: str = "Polarization[μC/cm^2]",
    limit_x: tuple[float, float] = None, 
    limit_y: tuple[float, float] = None,
) -> None:

    # 获取颜色映射
    colors = plt.get_cmap("tab20")(np.linspace(0, 1, 19))
    
    # 创建图形和坐标轴
    fig, ax = plt.subplots()
    plt.subplots_adjust(right=0.85)
    
    # 用于存储图例的句柄
    handles = []
    
    # 分别绘制每组散点
    for i in range(len(x)):
        ax.scatter(x[i], y[i], s=point_size, color=colors[i*3%19], label=group[i])
        handles.append(plt.Line2D([], [], color=colors[i*3%19], linestyle='-', label=group[i]))
    
    # 添加图例
    plt.legend(
        handles=handles,
        bbox_to_anchor=(1.00, 1),
        loc='upper left',
        borderaxespad=0,
        bbox_transform=plt.gca().transAxes
    )
    
    # 如果设置了 limit_x，则限制 x 轴范围
    if limit_x is not None:
        ax.set_xlim(limit_x)
    
    # 如果设置了 limit_y，则限制 y 轴范围
    if limit_y is not None:
        ax.set_ylim(limit_y)
    
    # 设置 x 轴和 y 轴的标签
    ax.set_xlabel(label_x)
    ax.set_ylabel(label_y)
    ax.set_title(title)
    
    # 显示图形
    plt.show()

if __name__ == "__main__":
    file_Name = get_info()
    x, y, group = read_file(file_Name)
    draw_plot(x, y, group)
